#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
by:杨锦成
vx:18873833517
time:2024.9.4
file_name:nav_position_pkg.py
说明:无人机参数MC_YAW_P=0.48
'''
import rospy
import numpy as np
from geometry_msgs.msg import Twist, PoseStamped
from mavros_msgs.msg import PositionTarget
from std_msgs.msg import Float32MultiArray
from filterpy.kalman import KalmanFilter
from nav_msgs.msg import Odometry
from actionlib_msgs.msg import GoalStatusArray
import threading
import math

class SmoothedControl:
    def __init__(self):
        self.velocity_x_kf = KalmanFilter(dim_x=2, dim_z=1)
        self.velocity_y_kf = KalmanFilter(dim_x=2, dim_z=1)
        self.yaw_rate_kf = KalmanFilter(dim_x=2, dim_z=1)

        self._initialize_kalman_filter(self.velocity_x_kf)
        self._initialize_kalman_filter(self.velocity_y_kf)
        self._initialize_kalman_filter(self.yaw_rate_kf)

    def _initialize_kalman_filter(self, kf):
        kf.x = np.array([0., 0.])
        kf.P = np.eye(2) * 1000.
        kf.F = np.array([[1., 1.],
                         [0., 1.]])
        kf.H = np.array([[1., 0.]])
        kf.R = np.array([[1]])
        kf.Q = np.array([[0.01, 0.],
                         [0.,   0.01]])

    def filter(self, kf, z):
        kf.predict()
        kf.update(z)
        return kf.x[0]

class HoverController:
    def __init__(self, action_mode="参数", log_function=rospy.loginfo):
        self.initial_yaw = 0.0
        self.initial_yaw_set = False  # 是否已经获取初始偏航角的标志
        self.goal_reached = False
        self.odom_recorded = False
        self.hover_position = None
        self.goal_status = None
        self.new_plan_detected = False
        self.odom_msg = None
        self.cmd_vel_active = True
        self.action_mode = action_mode
        self.log_function = log_function

        # 启动目标状态监听
        self.goal_thread = threading.Thread(target=self._goal_message_listener)
        self.goal_thread.daemon = True
        self.goal_thread.start()

        # 根据模式选择里程计订阅
        if self.action_mode in ["sim", "sim_nolog"]:
            self.odom_thread = threading.Thread(target=self._sim_odom_message_listener)
        else:
            self.odom_thread = threading.Thread(target=self._odom_message_listener)

        self.odom_thread.daemon = True
        self.odom_thread.start()

    def _goal_message_listener(self):
        rate = rospy.Rate(10)
        while not rospy.is_shutdown():
            try:
                self.goal_status = rospy.wait_for_message('/move_base/status', GoalStatusArray, timeout=0.5)
                self._check_for_new_plan(self.goal_status)
            except rospy.ROSException:
                continue
            rate.sleep()

    def _check_for_new_plan(self, status_msg):
        for status in status_msg.status_list:
            if status.status == 1:
                self.new_plan_detected = True
                self.log_function("检测到新路径规划，退出悬停模式，继续速度控制")
                return
        self.new_plan_detected = False

    def _odom_message_listener(self):
        self.log_function("订阅里程计话题：/mavros/local_position/odom")
        rospy.Subscriber('/mavros/local_position/odom', Odometry, self._odom_callback, queue_size=10)

    def _sim_odom_message_listener(self):
        self.log_function("订阅仿真里程计话题：/iris_0/mavros/local_position/odom")
        rospy.Subscriber('/iris_0/mavros/local_position/odom', Odometry, self._odom_callback, queue_size=10)

    def _odom_callback(self, msg):
        self.odom_msg = msg
        if not self.initial_yaw_set:
            self._update_initial_yaw()

    def _update_initial_yaw(self):
        """从里程计数据获取初始偏航角"""
        if self.odom_msg:
            orientation = self.odom_msg.pose.pose.orientation
            q = [orientation.x, orientation.y, orientation.z, orientation.w]
            self.initial_yaw = 0
            self.initial_yaw_set = True
            self.log_function(f"初始偏航角已设置：{math.degrees(self.initial_yaw):.2f}°")

    @staticmethod
    def quaternion_to_yaw(q):
        """四元数转偏航角计算"""
        return math.atan2(2.0*(q[3]*q[2] + q[0]*q[1]),
                          1.0 - 2.0*(q[1]**2 + q[2]**2))

    def process_goal_status(self):
        """判断目标是否到达"""
        if self.goal_status:
            for status in self.goal_status.status_list:
                if status.status == 3:
                    self.log_function("目标已到达，准备进入悬停模式")
                    self.goal_reached = True
                    break

    def process_new_goal(self):
        """判断是否检测到新的目标规划"""
        if self.new_plan_detected:
            self.goal_reached = False
            self.odom_recorded = False
            self.hover_position = None
            self.log_function("检测到新目标，重置悬停状态，继续速度控制")

    def record_odom(self):
        """在目标到达后记录当前位置"""
        if self.odom_msg:
            if not self.odom_recorded and self.goal_reached:
                self.hover_position = (
                    self.odom_msg.pose.pose.position.x,
                    self.odom_msg.pose.pose.position.y
                )
                self.odom_recorded = True
                self.log_function(f"已记录悬停位置：x={self.hover_position[0]}, y={self.hover_position[1]}")

    def get_current_altitude(self):
        """获取当前无人机高度"""
        if self.odom_msg:
            return self.odom_msg.pose.pose.position.z
        return None

def degrees_to_radians(degrees):
    return -degrees * (math.pi / 180.0)

def publish_nav_goal_x_y_z_yaw_t_frame(
    x, y, z, yaw, t, frame,
    setpoint_pub, nav_goal_pub,
    log_function=rospy.loginfo,
    action_mode="参数",
    stop_thread_flag=None
):
    log_function(f"开始发布导航目标... [模式={action_mode}]")

    rate = rospy.Rate(15)
    target = PositionTarget()

    smooth_control = SmoothedControl()
    hover_controller = HoverController(action_mode=action_mode, log_function=log_function)

    # 新增：用于存储从 /move_base_simple/goal 获取的目标坐标和偏航角
    goal_position = None
    goal_yaw = None

    def goal_callback(msg):
        nonlocal goal_position, goal_yaw
        # 提取目标位置
        goal_position = (msg.pose.position.x, msg.pose.position.y)
        # 提取目标偏航角（从四元数转换为 yaw）
        orientation = msg.pose.orientation
        q = [orientation.x, orientation.y, orientation.z, orientation.w]
        goal_yaw = hover_controller.quaternion_to_yaw(q)
        log_function(f"接收到新目标坐标：x={goal_position[0]}, y={goal_position[1]}, yaw={math.degrees(goal_yaw):.2f}°")

    # 订阅 /move_base_simple/goal 话题
    rospy.Subscriber('/move_base_simple/goal', PoseStamped, goal_callback)

    # 等待直到获取到初始偏航角
    while not hover_controller.initial_yaw_set and not rospy.is_shutdown():
        log_function("等待获取初始偏航角...")
        rospy.sleep(0.1)

    # 将发布的 yaw(角度)转换为相对于初始航向的弧度偏航量
    relative_yaw = hover_controller.initial_yaw + degrees_to_radians(yaw)

    # 先发布几次导航目标（兼容某些需要连续接收的场景）
    for i in range(3):
        nav_goal_msg = Float32MultiArray()
        nav_goal_msg.data = [x, y, relative_yaw, frame]
        nav_goal_pub.publish(nav_goal_msg)
        log_function("已发布导航目标")
        target.type_mask = 0b100111111000
        target.yaw = relative_yaw
        print(yaw)
        rospy.sleep(0.01)

    start_time = rospy.Time.now()
    end_time = start_time + rospy.Duration(t)

    # ---------------
    # 新增：统计进入“悬停模式条件”的次数
    # ---------------
    hover_mode_entry_count = 0

    while rospy.Time.now() < end_time:
        if rospy.is_shutdown():
            break
        if stop_thread_flag.is_set():
            rospy.logwarn("nav 线程被中断")
            break

        # 处理目标状态和新目标检测
        hover_controller.process_goal_status()
        hover_controller.process_new_goal()

        # 获取 /cmd_vel，如果没有收到则认为指令停止
        try:
            cmd_vel_msg = rospy.wait_for_message('/cmd_vel', Twist, timeout=0.1)
            hover_controller.cmd_vel_active = True
            target.type_mask = 0b110111000011

            # 使用 KalmanFilter 平滑速度和偏航率
            target.velocity.x = smooth_control.filter(smooth_control.velocity_x_kf, cmd_vel_msg.linear.x)
            target.velocity.y = smooth_control.filter(smooth_control.velocity_y_kf, cmd_vel_msg.linear.y)
            target.yaw_rate   = smooth_control.filter(smooth_control.yaw_rate_kf, cmd_vel_msg.angular.z)
            target.position.z = z

            # 始终维持 yaw 为初始角（+ 后续的相对偏航）
            target.yaw = relative_yaw
        except rospy.ROSException:
            hover_controller.cmd_vel_active = False

        # 获取当前高度
        current_altitude = hover_controller.get_current_altitude()

        # 如果 current_altitude 不为 None，则继续保持当前高度
        if current_altitude is not None:
            # 控制高度：通过计算当前高度与目标高度的差异来确定垂直速度
            altitude_error = z - current_altitude

            # 使用线性控制直接调整垂直速度
            vertical_velocity = np.clip(altitude_error * 3.0, -0.5, 0.5)
            target.velocity.z = vertical_velocity

        # 记录当前姿态用于悬停
        hover_controller.record_odom()

        # -----------------------
        # 下面是关键修改处：只有第三次进入该 if 才真正设置“悬停模式”
        # -----------------------
        if (not hover_controller.cmd_vel_active
            and hover_controller.goal_reached
            and hover_controller.odom_recorded
            and hover_controller.hover_position
        ):
            hover_mode_entry_count += 1  # 每次满足条件都 +1
            if hover_mode_entry_count >= 3:
                log_function("【第三次】速度指令停止且目标已到达，进入悬停模式")
                # 真正进入悬停模式
                target.type_mask = 0b100111111000

                # 如果接收到新的目标坐标和偏航角，则使用新值；否则使用默认值
                if goal_position:
                    target.position.x = goal_position[0]
                    target.position.y = goal_position[1]
                else:
                    target.position.x = x
                    target.position.y = y

                if goal_yaw is not None:
                    target.yaw = goal_yaw  # 使用新目标的偏航角
                else:
                    target.yaw = relative_yaw  # 使用默认的偏航角

                target.position.z = z  # 保持高度不变
                target.coordinate_frame = PositionTarget.FRAME_LOCAL_NED
                target.velocity.x = 0.0
                target.velocity.y = 0.0
            else:
                # 第 1、2 次（或超过 3 次之后的所有情况），
                # 不执行“悬停模式”的那一段控制指令，而是恢复到 LOCAL_NED 速度控制
                target.type_mask = 0b110111000011
                target.coordinate_frame = PositionTarget.FRAME_BODY_NED
        else:
            # 如果条件不满足，或者还没到第三次，就用原先的 LOCAL_NED 速度控制方式
            target.type_mask = 0b110111000011
            target.coordinate_frame = PositionTarget.FRAME_BODY_NED

        # 时间戳与发布
        target.header.stamp = rospy.Time.now()
        setpoint_pub.publish(target)

        # 打印调试信息
        log_function(
            "控制指令：vx=%.2f, vy=%.2f, vz=%.2f, yaw_rate=%.2f | "
            "位置(x=%.2f, y=%.2f, z=%.2f), yaw=%.2f | 第%d次满足条件"
            % (
                target.velocity.x, target.velocity.y, target.velocity.z, target.yaw_rate,
                target.position.x, target.position.y, target.position.z, target.yaw,
                hover_mode_entry_count
            )
        )

        rate.sleep()
